import numpy as np
from sklearn.decomposition import MiniBatchDictionaryLearning
import os
import timeit

PATCH_COUNT = 400000
COMP_COUNT = 350
CODE_COUNT = 2800
BATCH_SIZE = 500
EPOCH_COUNT = 16
ITERATION_COUNT = int(PATCH_COUNT / BATCH_SIZE) * EPOCH_COUNT
SPARSITY_PARAMETER = 4.0
NON_NEGATIVE = True
OUT_FOLDER = 'scbasis/'

start = timeit.default_timer()

if not os.path.exists(OUT_FOLDER):
    os.makedirs(OUT_FOLDER)

fileName = OUT_FOLDER + 'basisnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMP_COUNT) + 'codes' + str(CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'
print(fileName)

pcaV1C = np.load('imnet_inference/pcaTransformed.npy')[:, :COMP_COUNT]

dl = MiniBatchDictionaryLearning(n_components = CODE_COUNT, batch_size = BATCH_SIZE, alpha = SPARSITY_PARAMETER, transform_alpha = SPARSITY_PARAMETER, n_iter = ITERATION_COUNT, fit_algorithm = "cd", transform_algorithm = 'lasso_cd', positive_code = NON_NEGATIVE)

dl.fit(pcaV1C)

np.save(fileName, dl.components_)

stop = timeit.default_timer()

print('Time: ', stop - start)
